module ub_ram (
    input logic clk,
    input logic rst_n,
    ub_rd_req_if.in u_rd_req0_if,
    ub_rd_req_if.in u_rd_req1_if,
    gdma_ub_wr_req_if.in u_wr_req0_if,
    aru_ub_dat_if.out u_aru_normal_rd_dat_if,
    ldma_ub_dat_if.out u_aru_atomic_rd_dat_if,
    ldma_ub_dat_if.out u_ldma_rd_dat_if
);

    // UB RAM parameters (similar to ARB structure)
    localparam BANK_NUM = `N0 / `P_ARU;  // Number of banks
    localparam BANK_WIDTH = (`P_ARU * `N0) * 2;  // Data width per bank in bytes (32B per bank)
    localparam BANK_DEPTH = `UB_SIZE / (BANK_WIDTH * BANK_NUM);  // Number of entries per bank
    localparam ADDR_WIDTH = $clog2(BANK_DEPTH);

    // Internal signals for RAM control
    logic                        ram_wr_en   [BANK_NUM-1:0];
    logic                        ram_rd0_en  [BANK_NUM-1:0];
    logic                        ram_rd1_en  [BANK_NUM-1:0];
    logic     [  ADDR_WIDTH-1:0] ram_wr_addr [BANK_NUM-1:0];
    logic     [  ADDR_WIDTH-1:0] ram_rd0_addr[BANK_NUM-1:0];
    logic     [  ADDR_WIDTH-1:0] ram_rd1_addr[BANK_NUM-1:0];
    logic     [BANK_WIDTH*8-1:0] ram_wr_data [BANK_NUM-1:0];
    logic     [  BANK_WIDTH-1:0] ram_wr_mask [BANK_NUM-1:0];
    logic     [BANK_WIDTH*8-1:0] ram_rd0_data[BANK_NUM-1:0];
    logic     [BANK_WIDTH*8-1:0] ram_rd1_data[BANK_NUM-1:0];

    // Write request signals
    bf16_t                       wr_dat      [ `N0*`K0-1:0];  // GDMA writes full 128B (N0*K0 elements)
    logic                        wr_vld;
    ub_addr_t                    wr_addr;
    logic                        wr_msk      [ `N0*`K0-1:0];  // Unpacked array to match interface

    // Read request signals
    logic rd0_vld, rd1_vld;
    ub_addr_t rd0_addr, rd1_addr;
    logic [1:0] rd0_src, rd1_src;

    // Bank offset calculation
    logic [$clog2(BANK_NUM)-1:0] rd0_bank_offset, rd1_bank_offset, wr_bank_offset;
    assign rd0_bank_offset = rd0_addr[$clog2(BANK_WIDTH)];
    assign rd1_bank_offset = rd1_addr[$clog2(BANK_WIDTH)];
    assign wr_bank_offset  = wr_addr[$clog2(BANK_WIDTH)];

    //Depth offset calculation
    logic [$clog2(BANK_DEPTH)-1:0] rd0_depth_offset, rd1_depth_offset, wr_depth_offset;
    assign rd0_depth_offset = rd0_addr[$clog2(`UB_SIZE)-1 : $clog2(BANK_WIDTH*BANK_NUM)];
    assign rd1_depth_offset = rd1_addr[$clog2(`UB_SIZE)-1 : $clog2(BANK_WIDTH*BANK_NUM)];
    assign wr_depth_offset  = wr_addr[$clog2(`UB_SIZE)-1 : $clog2(BANK_WIDTH*BANK_NUM)];

    // Data preparation
    always_comb begin
        // Convert packed array to unpacked array
        for (int i = 0; i < `N0 * `K0; i++) begin
            wr_dat[i] = u_wr_req0_if.pld.dat[i];
            wr_msk[i] = u_wr_req0_if.pld.msk[i];
        end

        wr_vld   = u_wr_req0_if.vld;
        wr_addr  = u_wr_req0_if.pld.ub_addr;

        rd0_vld  = u_rd_req0_if.vld;
        rd0_addr = u_rd_req0_if.ub_addr;
        rd0_src  = u_rd_req0_if.src;

        rd1_vld  = u_rd_req1_if.vld;
        rd1_addr = u_rd_req1_if.ub_addr;
        rd1_src  = u_rd_req1_if.src;
    end

    // Generate BANK_NUM RAM banks
    genvar bank;
    generate
        for (bank = 0; bank < BANK_NUM; bank++) begin : gen_ram_banks
            common_ram_2r1w #(
                .ADDR_WIDTH(ADDR_WIDTH),
                .DATA_WIDTH(BANK_WIDTH),
                .DEPTH(BANK_DEPTH)
            ) u_bank_ram (
                .clk(clk),
                .wr_en(ram_wr_en[bank]),
                .wr_addr(ram_wr_addr[bank]),
                .wr_data(ram_wr_data[bank]),
                .wr_byte_mask(ram_wr_mask[bank]),
                .rd0_en(ram_rd0_en[bank]),
                .rd0_addr(ram_rd0_addr[bank]),
                .rd0_data(ram_rd0_data[bank]),
                .rd1_en(ram_rd1_en[bank]),
                .rd1_addr(ram_rd1_addr[bank]),
                .rd1_data(ram_rd1_data[bank])
            );
        end
    endgenerate

    // Write control logic
    always_comb begin
        // Default values
        for (int bank = 0; bank < BANK_NUM; bank++) begin : init_write_signals
            ram_wr_en[bank]   = 1'b0;
            ram_wr_addr[bank] = 'd0;
            ram_wr_data[bank] = 'd0;
            ram_wr_mask[bank] = 'd0;
        end

        if (wr_vld) begin
            // GDMA writes to all banks (N0*K0 elements = 128B)
            for (int bank = 0; bank < BANK_NUM; bank++) begin
                // Calculate logical bank index after rotation
                static int src_bank = 0;
                src_bank = (bank + wr_bank_offset) % BANK_NUM;

                // Pack data for this bank (P_ARU x K0 elements per bank)
                for (int p = 0; p < `N0; p++) begin
                    for (int k = 0; k < `K0; k++) begin
                        static int bit_pos = 0;
                        static int dat_idx = 0;
                        bit_pos = (p * `K0 + k) * 16;
                        dat_idx = bank * (`N0 * `K0) + p * `K0 + k;
                        if (dat_idx < `N0 * `K0) begin
                            ram_wr_data[src_bank][bit_pos+:16] = wr_dat[dat_idx];
                            ram_wr_addr[src_bank]              = wr_depth_offset;
                        end else begin
                            ram_wr_data[~src_bank][bit_pos+:16] = wr_dat[dat_idx%(`N0*`K0)];
                            ram_wr_addr[~src_bank]              = wr_depth_offset + ((wr_bank_offset + 1) % BANK_NUM);
                        end
                    end
                end

                // Pack mask for this bank
                for (int elem = 0; elem < (`N0 * `K0); elem++) begin
                    static int msk_idx = 0;
                    msk_idx = bank * (`N0 * `K0) + elem;
                    if (msk_idx < `N0 * `K0) begin
                        ram_wr_mask[src_bank][elem] = wr_msk[msk_idx];
                    end else begin
                        ram_wr_mask[~src_bank][elem] = wr_msk[msk_idx%(`N0*`K0)];
                    end
                end

                // Enable write if at least one mask bit is set
                ram_wr_en[bank] = |ram_wr_mask[bank];
            end
        end
    end

    // Read control logic
    always_comb begin
        // Default: disable all banks
        for (int bank = 0; bank < BANK_NUM; bank++) begin : init_read_signals
            ram_rd0_en[bank]   = 1'b0;
            ram_rd0_addr[bank] = 'd0;
            ram_rd1_en[bank]   = 1'b0;
            ram_rd1_addr[bank] = 'd0;
        end

        // Read port 0
        if (rd0_vld) begin
            if (rd0_src == 2'b10) begin  // LDMA: read all banks
                for (int bank = 0; bank < BANK_NUM; bank++) begin
                    ram_rd0_en[bank] = 1'b1;
                end
                ram_rd0_addr[rd0_bank_offset]  = rd0_depth_offset;
                ram_rd0_addr[~rd0_bank_offset] = rd0_depth_offset + ((rd0_bank_offset + 1) % BANK_NUM);
            end else begin  // ARU (normal/atomic): read single bank
                ram_rd0_en[rd0_bank_offset]   = 1'b1;
                ram_rd0_addr[rd0_bank_offset] = rd0_depth_offset;
            end
        end

        // Read port 1
        if (rd1_vld) begin
            if (rd1_src == 2'b10) begin  // LDMA: read all banks
                for (int bank = 0; bank < BANK_NUM; bank++) begin
                    ram_rd1_en[bank] = 1'b1;
                end
                ram_rd1_addr[rd1_bank_offset]  = rd1_depth_offset;
                ram_rd1_addr[~rd1_bank_offset] = rd1_depth_offset + ((rd1_bank_offset + 1) % BANK_NUM);
            end else begin  // ARU (normal/atomic): read single bank
                ram_rd1_en[rd1_bank_offset]   = 1'b1;
                ram_rd1_addr[rd1_bank_offset] = rd1_depth_offset;
            end
        end
    end

    // Pipeline registers for control signals
    // 声明用于存储控制信号的寄存器，'_q' 后缀通常表示这是寄存器的输出
    logic rd0_vld_q, rd1_vld_q;
    logic [1:0] rd0_src_q, rd1_src_q;
    logic [$clog2(BANK_NUM)-1:0] rd0_bank_offset_q, rd1_bank_offset_q;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            rd0_vld_q         <= 1'b0;
            rd0_src_q         <= 2'b00;
            rd0_bank_offset_q <= '0;
            rd1_vld_q         <= 1'b0;
            rd1_src_q         <= 2'b00;
            rd1_bank_offset_q <= '0;
        end else begin
            rd0_vld_q         <= rd0_vld;
            rd0_src_q         <= rd0_src;
            rd0_bank_offset_q <= rd0_bank_offset;
            rd1_vld_q         <= rd1_vld;
            rd1_src_q         <= rd1_src;
            rd1_bank_offset_q <= rd1_bank_offset;
        end
    end

    // Output data assembly with reordering
    // Use N0 x K0 format to accommodate both ARU (32B) and LDMA (64B) reads
    bf16_t rd0_dat[`N0-1:0][`K0-1:0];
    bf16_t rd1_dat[`N0-1:0][`K0-1:0];

    always_ff @(posedge clk) begin
        // Read port 0 data assembly
        if (rd0_src_q == 2'b10) begin  // LDMA: reorder data from all banks
            for (int bank_idx = 0; bank_idx < BANK_NUM; bank_idx++) begin
                for (int p = 0; p < `P_ARU; p++) begin
                    for (int k = 0; k < `K0; k++) begin
                        static int bit_pos = 0;
                        static int n_idx = 0;
                        bit_pos = (p * `K0 + k) * 16;
                        n_idx   = bank_idx * `P_ARU + p;
                        rd0_dat[n_idx][k] <= ram_rd0_data[(bank_idx+rd0_bank_offset_q)%BANK_NUM][bit_pos+:16];
                    end
                end
            end
        end else begin  // ARU (normal/atomic): only use data from target bank
            for (int p = 0; p < `P_ARU; p++) begin
                for (int k = 0; k < `K0; k++) begin
                    static int bit_pos = 0;
                    static int n_idx = 0;
                    bit_pos = (p * `K0 + k) * 16;
                    n_idx   = p;
                    rd0_dat[n_idx][k] <= ram_rd0_data[rd0_bank_offset_q][bit_pos+:16];
                end
            end
        end

        // Read port 1 data assembly
        if (rd1_src_q == 2'b10) begin  // LDMA: reorder data from all banks
            for (int bank_idx = 0; bank_idx < BANK_NUM; bank_idx++) begin
                for (int p = 0; p < `P_ARU; p++) begin
                    for (int k = 0; k < `K0; k++) begin
                        static int bit_pos = 0;
                        static int n_idx = 0;
                        bit_pos = (p * `K0 + k) * 16;
                        n_idx   = bank_idx * `P_ARU + p;
                        rd1_dat[n_idx][k] <= ram_rd1_data[(bank_idx+rd1_bank_offset_q)%BANK_NUM][bit_pos+:16];
                    end
                end
            end
        end else begin  // ARU (normal/atomic): only use data from target bank
            for (int p = 0; p < `P_ARU; p++) begin
                for (int k = 0; k < `K0; k++) begin
                    static int bit_pos = 0;
                    static int n_idx = 0;
                    bit_pos = (p * `K0 + k) * 16;
                    n_idx   = p;
                    rd1_dat[n_idx][k] <= ram_rd1_data[rd1_bank_offset_q][bit_pos+:16];
                end
            end
        end
    end

    // Output routing based on src signal
    logic rd0_to_aru_normal, rd0_to_aru_atomic, rd0_to_ldma;
    logic rd1_to_aru_normal, rd1_to_aru_atomic, rd1_to_ldma;

    assign rd0_to_aru_normal          = rd0_vld_q && (rd0_src_q == 2'b00);
    assign rd0_to_aru_atomic          = rd0_vld_q && (rd0_src_q == 2'b01);
    assign rd0_to_ldma                = rd0_vld_q && (rd0_src_q == 2'b10);

    assign rd1_to_aru_normal          = rd1_vld_q && (rd1_src_q == 2'b00);
    assign rd1_to_aru_atomic          = rd1_vld_q && (rd1_src_q == 2'b01);
    assign rd1_to_ldma                = rd1_vld_q && (rd1_src_q == 2'b10);

    // ARU normal read data output (P_ARU x N0 format = 32B)
    assign u_aru_normal_rd_dat_if.vld = rd0_to_aru_normal | rd1_to_aru_normal;
    always_comb begin
        for (int p = 0; p < `P_ARU; p++) begin
            for (int n = 0; n < `N0; n++) begin
                if (rd0_to_aru_normal) begin
                    u_aru_normal_rd_dat_if.dat[p][n] = rd0_dat[p][n];
                end else if (rd1_to_aru_normal) begin
                    u_aru_normal_rd_dat_if.dat[p][n] = rd1_dat[p][n];
                end else begin
                    u_aru_normal_rd_dat_if.dat[p][n] = '0;
                end
            end
        end
    end

    // ARU atomic read data output (N0 x K0 format = 128B, but only lower 64B valid)
    assign u_aru_atomic_rd_dat_if.vld = rd0_to_aru_atomic | rd1_to_aru_atomic;
    always_comb begin
        for (int n = 0; n < `N0; n++) begin
            for (int k = 0; k < `K0; k++) begin
                if (rd0_to_aru_atomic) begin
                    // Lower 64B (P_ARU rows): valid data from rd0_dat
                    if (n < `P_ARU) begin
                        u_aru_atomic_rd_dat_if.dat[n][k] = rd0_dat[n][k];
                    end else begin
                        // Upper 64B: fill with zeros
                        u_aru_atomic_rd_dat_if.dat[n][k] = '0;
                    end
                end else if (rd1_to_aru_atomic) begin
                    // Lower 64B (P_ARU rows): valid data from rd1_dat
                    if (n < `P_ARU) begin
                        u_aru_atomic_rd_dat_if.dat[n][k] = rd1_dat[n][k];
                    end else begin
                        // Upper 64B: fill with zeros
                        u_aru_atomic_rd_dat_if.dat[n][k] = '0;
                    end
                end else begin
                    u_aru_atomic_rd_dat_if.dat[n][k] = '0;
                end
            end
        end
    end

    // LDMA read data output (N0 x K0 format = 64B)
    assign u_ldma_rd_dat_if.vld = rd0_to_ldma | rd1_to_ldma;
    always_comb begin
        for (int n = 0; n < `N0; n++) begin
            for (int k = 0; k < `K0; k++) begin
                if (rd0_to_ldma) begin
                    u_ldma_rd_dat_if.dat[n][k] = rd0_dat[n][k];
                end else if (rd1_to_ldma) begin
                    u_ldma_rd_dat_if.dat[n][k] = rd1_dat[n][k];
                end else begin
                    u_ldma_rd_dat_if.dat[n][k] = '0;
                end
            end
        end
    end

    // RAM interfaces always ready (no backpressure)
    assign u_rd_req0_if.rdy = 1'b1;
    assign u_rd_req1_if.rdy = 1'b1;
    assign u_wr_req0_if.rdy = 1'b1;
    // assign u_aru_normal_rd_dat_if.rdy = 1'b1;
    // assign u_aru_atomic_rd_dat_if.rdy = 1'b1;
    // assign u_ldma_rd_dat_if.rdy = 1'b1;

endmodule
